# 在combination中已经训练好了样本，在这里进行读入的操作
# 导入相应的库
import keras.models
from tensorflow.keras.preprocessing.image import ImageDataGenerator
import matplotlib.pyplot as plt
import tensorflow as tf
from PIL import Image
import numpy as np
import os
import time


start = time.perf_counter()

# 设置图片的高和宽
im_height = 224
im_width = 224

#设置一次训练样本数和迭代次数
batch_size = 64    #一次训练所选取的样本数
epochs = 2   #训练时，使用所有数据集对模型进行一次完整的训练，称为一次epoch

#设置学习率，优化器，学习率下降有关的参数
# Init_lr = 1e-3
# Min_lr = Init_lr * 0.01

image_path = "../datasets/"  # 猫狗数据集路径
train_dir = "../datasets/train"  # 训练集路径
validation_dir = "../datasets/test"  # 验证集路径

# 定义训练集图像生成器，并进行图像增强
# ImageDataGenerator通过实时数据产生一系列的旋转、平移、剪切、反转等变换过的图像
train_image_generator = ImageDataGenerator(rescale=1. / 255,  # 归一化
                                           rotation_range=40,  # 旋转范围
                                           width_shift_range=0.2,  # 水平平移范围
                                           height_shift_range=0.2,  # 垂直平移范围
                                           shear_range=0.2,  # 剪切变换的程度
                                           zoom_range=0.2,  # 剪切变换的程度
                                           horizontal_flip=True,  # 水平翻转
                                           fill_mode='nearest')

# 使用图像生成器从文件夹train_dir中读取样本，对标签进行one-hot编码
# 为什么要用one-hot编码？因为猫狗之间是无序的，如果用1、2就会有大小的区别，
# 故使用01 10式的编码，使其只有一位有效
train_data_gen = train_image_generator.flow_from_directory(directory=train_dir,  # 从训练集路径读取图片
                                                           batch_size=batch_size,  # 一次训练所选取的样本数
                                                           shuffle=True,  # 打乱标签
                                                           target_size=(im_height, im_width),  # 图片resize到224x224大小
                                                           class_mode='categorical')  # one-hot编码

# 训练集样本数
total_train = train_data_gen.n

# 定义验证集图像生成器，并对图像进行预处理
validation_image_generator = ImageDataGenerator(rescale=1. / 255)  # 归一化

# 使用图像生成器从验证集validation_dir中读取样本
val_data_gen = validation_image_generator.flow_from_directory(directory=validation_dir,  # 从验证集路径读取图片
                                                              batch_size=batch_size,  # 一次训练所选取的样本数
                                                              shuffle=False,  # 不打乱标签
                                                              target_size=(im_height, im_width),  # 图片resize到224x224大小
                                                              class_mode='categorical')  # one-hot编码

# 验证集样本数
total_val = val_data_gen.n

# 测试读取图片用时
end = time.perf_counter()
print('读取图片用时: %s Seconds' % (end - start))

start = time.perf_counter()
# model =ResNet50(input_shape=[224,224,3], classes=2)
# model.load_weights('./save_weights/ResNet50_tf.h5')
model = keras.models.load_model('./save_weights/ResNet50_tf.h5', compile=False) #预训模型基础上训练的
# model = keras.models.load_model('./save_weights/ResNet50.h5', compile=False) #从0开始训练的

end = time.perf_counter()
print('加载用时: %s Seconds' % (end - start))
#  获取数据集的类别编码
class_indices = train_data_gen.class_indices
# 将编码和对应的类别存入字典
inverse_dict = dict((val, key) for key, val in class_indices.items())
# 加载测试图片
# img = Image.open("../datasets/test/2.jpg")
# img = Image.open("../test1.png")
# img = Image.open("../test2.png")
img = Image.open("../test5.jpeg")
# img = Image.open("../cat.10000.jpg")

print("已加载图片")
# 将图片resize到250x250大小
img = img.resize((224, 224))
# 归一化
img1 = np.array(img) / 255.
# 将图片增加一个维度，目的是匹配网络模型
img1 = (np.expand_dims(img1, 0))
# 将预测结果转化为概率值
result = np.squeeze(model.predict(img1))
predict_class = np.argmax(result)
print(inverse_dict[int(predict_class)], result[predict_class])
# 将预测的结果打印在图片上面
plt.title([inverse_dict[int(predict_class)], result[predict_class]])
# 显示图片
plt.imshow(img)
plt.show()
